import inspect
import os
from pathlib import Path
import imgaug.augmenters as iaa

from text_renderer.effect import *
from text_renderer.corpus import *
from text_renderer.config import (
    RenderCfg,
    NormPerspectiveTransformCfg,
    GeneratorCfg,
    FixedTextColorCfg, SimpleTextColorCfg,
)
from text_renderer.layout.same_line import SameLineLayout
from text_renderer.layout.extra_text_line import ExtraTextLineLayout

CURRENT_DIR = Path(os.path.abspath(os.path.dirname(__file__)))
READER_ROOT = Path('D:/ocr/ocr_resources/reader/')

CONFIG_ROOT = Path('D:/workspace/python/train_data/train_data_configs/')

RENDER_ROOT = Path('D:/workspace/python/text_renderer/')

font_cfg = dict(
    font_dir=CONFIG_ROOT / 'fonts/chn',
    font_list_file=CONFIG_ROOT / "fonts/chn_windows.txt",
    font_size=(30, 31),
)

perspective_transform = NormPerspectiveTransformCfg(20, 20, 1.5)


def base_cfg(
        name: str, corpus, corpus_effects=None, layout_effects=None, layout=None, gray=True
):
    return GeneratorCfg(
        # num_image=6500000,
        # num_image=500000,
        num_image=100,
        save_dir=READER_ROOT / 'train_data' / 'train_data_zy650w_7',
        render_cfg=RenderCfg(
            bg_dir=RENDER_ROOT / 'example_data/bg',
            perspective_transform=perspective_transform,
            gray=gray,
            layout_effects=layout_effects,
            layout=layout,
            corpus=corpus,
            corpus_effects=corpus_effects,
        ),
    )


def train_data_zy650w():
    return base_cfg(
        inspect.currentframe().f_code.co_name,
        layout=ExtraTextLineLayout(),
        # layout=SameLineLayout(),
        # 图像是否变灰色,失去颜色
        gray=True,
        # gray=False,
        corpus=[EnumCorpus(
            EnumCorpusCfg(
                text_paths=[READER_ROOT / "all_words/in_chars_words_clear.txt"],
                # text_paths=[READER_ROOT / "all_words/demo.txt"],
                # num_pick=6040126,
                # text_paths (List[Path]): Text file paths
                # items (List[str]): Texts to choice. Only works if text_paths is empty
                # num_pick (int): Random choice {count} item from texts
                # filter_font (bool): Only work when filter_by_chars is True. If True, filter font file
                # by intersection of font support chars with chars file
                filter_by_chars=True,
                filter_font=True,
                chars_file=CONFIG_ROOT / "words/result_all.txt",
                **font_cfg
            ),
        )],
        corpus_effects=Effects([Line(p=0.01)]),
        # corpus_effects=[Effects([]), NoEffects() ],
        # corpus_effects=Effects(
        #     [Padding(),
        #      Line(0.1, color_cfg=FixedTextColorCfg()),
        #      OneOf([DropoutRand(), DropoutVertical()]),
        #      Padding(p=1, w_ratio=[0.2, 0.21], h_ratio=[0.7, 0.71], center=True),
        #      ImgAugEffect(aug=iaa.Emboss(alpha=(0.9, 1.0), strength=(1.5, 1.6))),
        #      MotionBlur()
        #      ]
        # ),
        # corpus_effects =    NoEffects(),
        # layout_effects=Effects(Line(p=0.3)),
    )


def train_data():
    return base_cfg(
        inspect.currentframe().f_code.co_name,
        layout=ExtraTextLineLayout(),
        # layout=SameLineLayout(),
        # 图像是否变灰色,失去颜色
        gray=True,
        # gray=False,
        corpus=[EnumCorpus(
            EnumCorpusCfg(
                # text_paths=[READER_ROOT / "all_words/in_chars_words_clear.txt"],
                # text_paths=[READER_ROOT / "all_words/demo.txt"],
                # num_pick=6040126,
                # text_paths (List[Path]): Text file paths
                # items (List[str]): Texts to choice. Only works if text_paths is empty
                # num_pick (int): Random choice {count} item from texts
                # filter_font (bool): Only work when filter_by_chars is True. If True, filter font file
                # by intersection of font support chars with chars file
                filter_by_chars=False,
                filter_font=False,
                auto_load_texts=False,
                chars_file=CONFIG_ROOT / "words/result_all.txt",
                **font_cfg
            ),
        )],
        corpus_effects=Effects([Line(p=0.01)]),
        # corpus_effects=[Effects([]), NoEffects() ],
        # corpus_effects=Effects(
        #     [Padding(),
        #      Line(0.1, color_cfg=FixedTextColorCfg()),
        #      OneOf([DropoutRand(), DropoutVertical()]),
        #      Padding(p=1, w_ratio=[0.2, 0.21], h_ratio=[0.7, 0.71], center=True),
        #      ImgAugEffect(aug=iaa.Emboss(alpha=(0.9, 1.0), strength=(1.5, 1.6))),
        #      MotionBlur()
        #      ]
        # ),
        # corpus_effects =    NoEffects(),
        # layout_effects=Effects(Line(p=0.3)),
    )


# fmt: off
# The configuration file must have a configs variable
configs = [
    # chn_data(),
    # train_data_zy650w(),
    train_data()
]
# fmt: on
